#include "kernel_operator.h"
#include <type_traits>
#include <float.h>
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;

// min max != 0
template<typename TYPE_X, typename TYPE_Y>
class KernelHistogram {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    TBuf<TPosition::VECCALC> xBuf;

    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_Y> yGm;

    uint32_t ALIGN_NUM;
    uint32_t totalLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;

    int32_t bins;
    float min;
    float max;
    float bin_width;

public:
    __aicore__ inline KernelHistogram() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t ALIGN_NUM, uint32_t totalLength, uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength, int32_t bins, float min , float max)
    {

        this->ALIGN_NUM = ALIGN_NUM;
        this->totalLength = totalLength;
        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;
        this->bins = bins;
        this->min = min;
        this->max = max;
        this->bin_width = (max - min) / bins;

        xGm.SetGlobalBuffer((__gm__ TYPE_X *)x, totalLength);
        yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y, bins);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));
        
        if constexpr(std::is_same_v<TYPE_X, int32_t> || std::is_same_v<TYPE_X, half>)    
            pipe.InitBuffer(xBuf, BUFFER_NUM * this->tileLength * sizeof(float));
        

    }

    // max, min为给定值
    __aicore__ inline void Process()
    {

        
        this->bin_width = (max - min) / bins;

        uint32_t loopCount = tileNum;
        for(int i=0; i<loopCount-1; i++)
        {
            CopyIn(i, tileLength);
            ComputeBins(i, tileLength);
        }
            CopyIn(tileNum-1, lasttileLength);
            ComputeBins(tileNum-1, totalLength - (tileNum - 1) * tileLength);
    }

    __aicore__ inline void CopyIn(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        DataCopy(xLocal, xGm[progress * tileLength], length);
        inQueueX.EnQue<TYPE_X>(xLocal);
    }

    __aicore__ inline void ComputeBins(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        if constexpr(std::is_same_v<TYPE_X, int32_t> || std::is_same_v<TYPE_X, half>)
        {
            LocalTensor<float> xLocal_f = xBuf.Get<float>();
            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);
            for (int i = 0; i < length; i++)
            {
                float var = xLocal_f.GetValue(i);
                if(var < min || var > max)
                    continue;
                int bin = (var - min) / bin_width;
                if(var == max)  //left close right open
                {
                    bin -= 1;
                }
                if constexpr(std::is_same_v<TYPE_X, half>)
                    yGm.SetValue(bin, (half)((float)yGm.GetValue(bin) + (float)1.0));
                else
                    yGm.SetValue(bin, yGm.GetValue(bin) + (TYPE_X)1);
            }
            
        }
        else    //float
        {
            for (int i = 0; i < length; i++)
            {
                float var = xLocal.GetValue(i);
                if(var < min || var > max)
                    continue;
                int bin = (var - min) / bin_width;
                if(var == max)  //left close right open
                {
                    bin -= 1;
                }
                
                yGm.SetValue(bin, yGm.GetValue(bin) + (float)1);
            }
        }
        
        inQueueX.FreeTensor<TYPE_X>(xLocal);
    }

};


//min max == 0
template<typename TYPE_X, typename TYPE_Y>
class KernelHistogramV1 {
private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    // TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;

    GlobalTensor<TYPE_X> xGm;
    GlobalTensor<TYPE_Y> yGm;
    TBuf<TPosition::VECCALC> Tmp, calcBuf, xBuf;

    uint32_t minLength;

    uint32_t ALIGN_NUM;
    uint32_t totalLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lasttileLength;

    int32_t bins;
    float min;
    float max;
    float bin_width;

public:
    __aicore__ inline KernelHistogramV1() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t ALIGN_NUM, uint32_t totalLength, uint32_t tileNum, uint32_t tileLength, uint32_t lasttileLength, int32_t bins)
    {

        this->ALIGN_NUM = ALIGN_NUM;
        this->minLength = 32/ALIGN_NUM;
        this->totalLength = totalLength;
        this->tileNum = tileNum;
        this->tileLength = tileLength;
        this->lasttileLength = lasttileLength;
        this->bins = bins;
        //适应float类型的最值
        this->min = FLT_MAX;
        this->max = -1 * FLT_MAX;
        
        xGm.SetGlobalBuffer((__gm__ TYPE_X *)x, totalLength);
        yGm.SetGlobalBuffer((__gm__ TYPE_Y *)y, bins);
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));
        // pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(TYPE_Y));
        
        if constexpr(std::is_same_v<TYPE_X, int32_t> || std::is_same_v<TYPE_X, half>)
        {
            pipe.InitBuffer(Tmp, BUFFER_NUM * this->minLength * sizeof(float));
            pipe.InitBuffer(xBuf, BUFFER_NUM * this->tileLength * sizeof(float));
            pipe.InitBuffer(calcBuf, BUFFER_NUM * this->tileLength * sizeof(float));
        }
        else    // float
        {
            pipe.InitBuffer(Tmp, BUFFER_NUM * this->minLength * sizeof(TYPE_X));
            // pipe.InitBuffer(xBuf, BUFFER_NUM * this->tileLength * sizeof(TYPE_X));
            pipe.InitBuffer(calcBuf, BUFFER_NUM * this->tileLength * sizeof(TYPE_X));
        }
        

    }

    __aicore__ inline void findMaxMin()
    {
        uint32_t loopCount = tileNum;
        for(int i=0; i<loopCount-1; i++)
        {
            CopyIn(i, tileLength);
            Compute(i, tileLength);
        }

        CopyIn(tileNum-1, lasttileLength);
        Compute(tileNum-1, totalLength - (tileNum - 1) * tileLength);
    }

    __aicore__ inline void CopyIn(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        DataCopy(xLocal, xGm[progress * tileLength], length);
        inQueueX.EnQue<TYPE_X>(xLocal);
    }

    __aicore__ inline void Compute(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        if constexpr(std::is_same_v<TYPE_X, int32_t> || std::is_same_v<TYPE_X, half>)   //int32先转为float才能计算reduce max
        {
            LocalTensor<float> wLocal = calcBuf.Get<float>();
            LocalTensor<float> tLocal = Tmp.Get<float>();
            LocalTensor<float> xLocal_f = xBuf.Get<float>();
            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);
            ReduceMax(tLocal, xLocal_f, wLocal, length);
            float var = tLocal.GetValue(0);
            // printf("maxop var: %f \n", var);
            if(var > max)
                max = var;
            ReduceMin(tLocal, xLocal_f, wLocal, length);    
            var = tLocal.GetValue(0);
            // printf("maxop var: %f \n", var);
            if(var < min)
                min = var;
        }
        else
        {
            LocalTensor<TYPE_X> wLocal = calcBuf.Get<TYPE_X>();
            LocalTensor<TYPE_X> tLocal = Tmp.Get<TYPE_X>();
            
            ReduceMax(tLocal, xLocal, wLocal, length);
            float var = tLocal.GetValue(0);
            // printf("maxop var: %f \n", var);
            if(var > max)
                max = var;
            ReduceMin(tLocal, xLocal, wLocal, length);    
            var = tLocal.GetValue(0);
            // printf("maxop var: %f \n", var);
            if(var < min)
                min = var;
        }
        
        inQueueX.FreeTensor<TYPE_X>(xLocal);
    }

    __aicore__ inline void ComputeBins(int32_t progress, uint32_t length)
    {
        LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        if constexpr(std::is_same_v<TYPE_X, int32_t> || std::is_same_v<TYPE_X, half>)
        {

            LocalTensor<float> xLocal_f = xBuf.Get<float>();
            Cast(xLocal_f, xLocal, RoundMode::CAST_NONE, length);
            // Muls(xLocal_f, xLocal_f, min, length);

            for (int i = 0; i < length; i++)
            {
                float var = xLocal_f.GetValue(i);
                if(var < min || var > max)
                    continue;
                int bin = (var - min) / bin_width;
                if(var == max)  //left close right open
                {
                    bin -= 1;
                }
                
                if constexpr(std::is_same_v<TYPE_X, half>)
                    yGm.SetValue(bin, (half)((float)yGm.GetValue(bin) + (float)1));
                else
                    yGm.SetValue(bin, yGm.GetValue(bin) + (TYPE_X)1);
            }
            
        }
        else    //float
        {
            for (int i = 0; i < length; i++)
            {
                float var = xLocal.GetValue(i);
                if(var < min || var > max)
                    continue;
                int bin = (var - min) / bin_width;
                if(var == max)  //left close right open
                {
                    bin -= 1;
                }
                
                yGm.SetValue(bin, yGm.GetValue(bin) + (float)1);
            }
        }
        
        inQueueX.FreeTensor<TYPE_X>(xLocal);
    }

    // max, min为给定值
    __aicore__ inline void Process()
    {
        findMaxMin();
        this->bin_width = (max - min) / bins;
        uint32_t loopCount = tileNum;
        for(int i=0; i<loopCount-1; i++)
        {
            CopyIn(i, tileLength);
            ComputeBins(i, tileLength);
        }
            CopyIn(tileNum-1, lasttileLength);
            ComputeBins(tileNum-1, totalLength - (tileNum - 1) * tileLength);
    }

};


extern "C" __global__ __aicore__ void histogram(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    // TODO: user kernel impl
    if(TILING_KEY_IS(1))
    {
        // printf("TILING_KEY_IS: 1 \n");
        KernelHistogramV1<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, tiling_data.ALIGN_NUM, tiling_data.totalLength, tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileLength, tiling_data.bins);
        op.Process();
    }
    else if (TILING_KEY_IS(2))
    {
        // printf("TILING_KEY_IS: 2 \n");
        KernelHistogram<DTYPE_X, DTYPE_Y> op;
        op.Init(x, y, tiling_data.ALIGN_NUM, tiling_data.totalLength, tiling_data.tileNum, tiling_data.tileLength, tiling_data.lasttileLength, tiling_data.bins, tiling_data.min , tiling_data.max);
        op.Process();
    }
    
}